"""
智能体注册管理器

负责智能体的注册、发现、管理和协调。
"""

from typing import Dict, List, Optional, Type, Any
from datetime import datetime
import asyncio
from loguru import logger

from .base_agent import BaseAgent, Task, Message, AgentStatus
from .agent_roles import AgentRole, AgentLevel, get_agent_config


class AgentRegistry:
    """智能体注册管理器"""
    
    def __init__(self):
        """初始化注册管理器"""
        self.agents: Dict[str, BaseAgent] = {}  # agent_id -> agent
        self.agents_by_role: Dict[AgentRole, List[BaseAgent]] = {}  # role -> agents
        self.agents_by_level: Dict[AgentLevel, List[BaseAgent]] = {}  # level -> agents
        self.agent_classes: Dict[AgentRole, Type[BaseAgent]] = {}  # role -> agent_class
        
        # 消息路由
        self.message_bus: List[Message] = []
        self.message_subscribers: Dict[str, List[str]] = {}  # message_type -> agent_ids
        
        # 任务分发
        self.global_task_queue: List[Task] = []
        self.task_assignments: Dict[str, str] = {}  # task_id -> agent_id
        
        # 性能监控
        self.system_metrics = {
            "total_agents": 0,
            "active_agents": 0,
            "total_tasks": 0,
            "completed_tasks": 0,
            "failed_tasks": 0,
            "average_response_time": 0.0,
            "system_load": 0.0,
            "last_update": datetime.now()
        }
        
        logger.info("智能体注册管理器初始化完成")
    
    async def initialize(self):
        """异步初始化智能体注册管理器"""
        try:
            # 初始化系统指标
            self._update_system_metrics()
            
            # 清理全局任务队列
            self.global_task_queue.clear()
            
            # 重置任务分配记录
            self.task_assignments.clear()
            
            logger.info("智能体注册管理器异步初始化完成")
            
        except Exception as e:
            logger.error(f"智能体注册管理器初始化失败: {e}")
            raise
    
    async def close(self):
        """关闭智能体注册管理器"""
        try:
            # 停止所有智能体
            for agent_id in list(self.agents.keys()):
                await self.unregister_agent(agent_id)
            
            # 清理资源
            self.agents.clear()
            self.agents_by_role.clear()
            self.agents_by_level.clear()
            self.message_bus.clear()
            self.message_subscribers.clear()
            self.global_task_queue.clear()
            self.task_assignments.clear()
            
            logger.info("智能体注册管理器已关闭")
            
        except Exception as e:
            logger.error(f"关闭智能体注册管理器失败: {e}")
            raise
    
    def register_agent_class(self, role: AgentRole, agent_class: Type[BaseAgent]):
        """注册智能体类"""
        self.agent_classes[role] = agent_class
        logger.info(f"注册智能体类: {role.value} -> {agent_class.__name__}")
    
    async def create_agent(
        self,
        role: AgentRole,
        agent_id: Optional[str] = None,
        name: Optional[str] = None,
        config: Optional[Dict[str, Any]] = None
    ) -> Optional[BaseAgent]:
        """创建智能体实例"""
        try:
            # 获取智能体类
            agent_class = self.agent_classes.get(role)
            if not agent_class:
                logger.error(f"未找到角色 {role.value} 对应的智能体类")
                return None
            
            # 生成agent_id
            if not agent_id:
                agent_id = f"{role.value}_{len(self.agents) + 1}"
            
            # 检查ID是否已存在
            if agent_id in self.agents:
                logger.error(f"智能体ID {agent_id} 已存在")
                return None
            
            # 创建智能体实例
            agent = agent_class(
                agent_id=agent_id,
                role=role,
                name=name,
                config=config
            )
            
            # 注册智能体
            await self.register_agent(agent)
            
            logger.info(f"创建智能体成功: {agent.name} ({agent_id})")
            return agent
            
        except Exception as e:
            logger.error(f"创建智能体失败: {e}")
            return None
    
    async def register_agent(self, agent: BaseAgent):
        """注册智能体"""
        try:
            # 添加到主注册表
            self.agents[agent.agent_id] = agent
            
            # 按角色分组
            if agent.role not in self.agents_by_role:
                self.agents_by_role[agent.role] = []
            self.agents_by_role[agent.role].append(agent)
            
            # 按层级分组
            if agent.level not in self.agents_by_level:
                self.agents_by_level[agent.level] = []
            self.agents_by_level[agent.level].append(agent)
            
            # 启动智能体
            await agent.start()
            
            # 更新系统指标
            self._update_system_metrics()
            
            logger.info(f"智能体 {agent.name} 注册成功")
            
        except Exception as e:
            logger.error(f"注册智能体失败: {e}")
            raise
    
    async def unregister_agent(self, agent_id: str):
        """注销智能体"""
        try:
            agent = self.agents.get(agent_id)
            if not agent:
                logger.warning(f"智能体 {agent_id} 不存在")
                return
            
            # 停止智能体
            await agent.stop()
            
            # 从注册表中移除
            del self.agents[agent_id]
            
            # 从角色分组中移除
            if agent.role in self.agents_by_role:
                self.agents_by_role[agent.role] = [
                    a for a in self.agents_by_role[agent.role] 
                    if a.agent_id != agent_id
                ]
            
            # 从层级分组中移除
            if agent.level in self.agents_by_level:
                self.agents_by_level[agent.level] = [
                    a for a in self.agents_by_level[agent.level]
                    if a.agent_id != agent_id
                ]
            
            # 更新系统指标
            self._update_system_metrics()
            
            logger.info(f"智能体 {agent.name} 注销成功")
            
        except Exception as e:
            logger.error(f"注销智能体失败: {e}")
            raise
    
    def get_agent(self, agent_id: str) -> Optional[BaseAgent]:
        """获取智能体"""
        return self.agents.get(agent_id)
    
    def get_agents_by_role(self, role: AgentRole) -> List[BaseAgent]:
        """根据角色获取智能体列表"""
        return self.agents_by_role.get(role, [])
    
    def get_agents_by_level(self, level: AgentLevel) -> List[BaseAgent]:
        """根据层级获取智能体列表"""
        return self.agents_by_level.get(level, [])
    
    def get_available_agents(self, role: Optional[AgentRole] = None) -> List[BaseAgent]:
        """获取可用的智能体列表"""
        agents = []
        
        if role:
            candidates = self.get_agents_by_role(role)
        else:
            candidates = list(self.agents.values())
        
        for agent in candidates:
            if agent.is_available:
                agents.append(agent)
        
        return agents
    
    async def assign_task(
        self,
        task: Task,
        target_role: Optional[AgentRole] = None,
        target_agent_id: Optional[str] = None
    ) -> bool:
        """分配任务"""
        try:
            # 指定智能体
            if target_agent_id:
                agent = self.get_agent(target_agent_id)
                if agent and agent.is_available:
                    success = await agent.assign_task(task)
                    if success:
                        self.task_assignments[task.id] = agent.agent_id
                        self._update_task_metrics()
                    return success
                return False
            
            # 指定角色
            if target_role:
                available_agents = self.get_available_agents(target_role)
            else:
                available_agents = self.get_available_agents()
            
            if not available_agents:
                # 添加到全局任务队列
                self.global_task_queue.append(task)
                logger.info(f"任务 {task.name} 添加到全局队列等待处理")
                return True
            
            # 选择最合适的智能体（简单的负载均衡）
            best_agent = min(available_agents, key=lambda a: len(a.task_queue))
            
            success = await best_agent.assign_task(task)
            if success:
                self.task_assignments[task.id] = best_agent.agent_id
                self._update_task_metrics()
            
            return success
            
        except Exception as e:
            logger.error(f"分配任务失败: {e}")
            return False
    
    async def broadcast_message(
        self,
        message: Message,
        target_level: Optional[AgentLevel] = None,
        target_role: Optional[AgentRole] = None
    ):
        """广播消息"""
        try:
            targets = []
            
            if target_role:
                targets = self.get_agents_by_role(target_role)
            elif target_level:
                targets = self.get_agents_by_level(target_level)
            else:
                targets = list(self.agents.values())
            
            for agent in targets:
                await agent.receive_message(message)
            
            logger.info(f"广播消息给 {len(targets)} 个智能体")
            
        except Exception as e:
            logger.error(f"广播消息失败: {e}")
    
    async def route_message(self, message: Message):
        """路由消息"""
        try:
            # 点对点消息
            if message.receiver and message.receiver != "broadcast":
                agent = self.get_agent(message.receiver)
                if agent:
                    await agent.receive_message(message)
                    return
            
            # 广播消息
            if message.receiver == "broadcast":
                await self.broadcast_message(message)
                return
            
            # 基于消息类型的订阅路由
            subscribers = self.message_subscribers.get(message.message_type, [])
            for agent_id in subscribers:
                agent = self.get_agent(agent_id)
                if agent:
                    await agent.receive_message(message)
            
        except Exception as e:
            logger.error(f"路由消息失败: {e}")
    
    def subscribe_message(self, agent_id: str, message_type: str):
        """订阅消息类型"""
        if message_type not in self.message_subscribers:
            self.message_subscribers[message_type] = []
        
        if agent_id not in self.message_subscribers[message_type]:
            self.message_subscribers[message_type].append(agent_id)
    
    def unsubscribe_message(self, agent_id: str, message_type: str):
        """取消订阅消息类型"""
        if message_type in self.message_subscribers:
            self.message_subscribers[message_type] = [
                aid for aid in self.message_subscribers[message_type]
                if aid != agent_id
            ]
    
    def _update_system_metrics(self):
        """更新系统指标"""
        self.system_metrics["total_agents"] = len(self.agents)
        self.system_metrics["active_agents"] = len([
            a for a in self.agents.values()
            if a.status != AgentStatus.OFFLINE
        ])
        
        # 计算系统负载
        total_queue_size = sum(len(a.task_queue) for a in self.agents.values())
        self.system_metrics["system_load"] = total_queue_size / max(len(self.agents), 1)
        
        self.system_metrics["last_update"] = datetime.now()
    
    def _update_task_metrics(self):
        """更新任务指标"""
        total_completed = sum(a.metrics["tasks_completed"] for a in self.agents.values())
        total_failed = sum(a.metrics["tasks_failed"] for a in self.agents.values())
        
        self.system_metrics["completed_tasks"] = total_completed
        self.system_metrics["failed_tasks"] = total_failed
        self.system_metrics["total_tasks"] = total_completed + total_failed
        
        # 计算平均响应时间
        response_times = [
            a.metrics["average_response_time"] for a in self.agents.values()
            if a.metrics["average_response_time"] > 0
        ]
        
        if response_times:
            self.system_metrics["average_response_time"] = sum(response_times) / len(response_times)
    
    def get_system_status(self) -> Dict[str, Any]:
        """获取系统状态"""
        self._update_system_metrics()
        self._update_task_metrics()
        
        return {
            "metrics": self.system_metrics,
            "agents": {
                agent_id: agent.get_status_info()
                for agent_id, agent in self.agents.items()
            },
            "global_queue_size": len(self.global_task_queue),
            "message_subscribers": self.message_subscribers
        }


# 全局智能体注册管理器实例
agent_registry = AgentRegistry()